/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.lvyh.lightframe.tracer.starter.kafka;

import com.lvyh.lightframe.tracer.core.common.AbstractTracer;
import com.lvyh.lightframe.tracer.core.common.SimpleTracerConfiguration;
import com.lvyh.lightframe.tracer.core.samplers.TraceSamplerConfig;
import com.lvyh.lightframe.tracer.plugin.kafka.TracerKafkaConsumerFactory;
import com.lvyh.lightframe.tracer.plugin.kafka.TracerKafkaProducerFactory;
import com.lvyh.lightframe.tracer.plugin.kafka.aop.KafkaListenerTracerAspect;
import com.lvyh.lightframe.tracer.plugin.kafka.tracers.KafkaConsumeTracer;
import com.lvyh.lightframe.tracer.plugin.kafka.tracers.KafkaSendTracer;
import com.lvyh.lightframe.tracer.starter.registry.DynamicBeanRegister;
import com.lvyh.lightframe.tracer.starter.util.SpringContextUtils;
import org.apache.kafka.clients.consumer.ConsumerConfig;
import org.apache.kafka.clients.producer.ProducerConfig;
import org.apache.kafka.common.serialization.StringDeserializer;
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.context.EnvironmentAware;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Import;
import org.springframework.core.env.Environment;
import org.springframework.kafka.annotation.EnableKafka;
import org.springframework.kafka.config.ConcurrentKafkaListenerContainerFactory;
import org.springframework.kafka.config.KafkaListenerContainerFactory;
import org.springframework.kafka.core.ConsumerFactory;
import org.springframework.kafka.core.DefaultKafkaConsumerFactory;
import org.springframework.kafka.core.DefaultKafkaProducerFactory;
import org.springframework.kafka.core.KafkaTemplate;
import org.springframework.kafka.core.ProducerFactory;
import org.springframework.kafka.listener.ConcurrentMessageListenerContainer;

import java.util.HashMap;
import java.util.Map;

@Configuration
@Import({DynamicBeanRegister.class})
@EnableKafka
@ConditionalOnProperty(prefix = "tracer.kafka", value = "enabled", havingValue = "true")
public class TracerKafkaAutoConfiguration implements EnvironmentAware {

    /**
     * Configuration context is the tool to get configuration files
     */
    private Environment env;

    @Bean
    public ProducerFactory<Integer, String> producerFactory() {
        TraceSamplerConfig traceSamplerConfig = SpringContextUtils.getBean(TraceSamplerConfig.class);
        AbstractTracer kafkaSendTracer = KafkaSendTracer.getKafkaSendTracer(traceSamplerConfig);

        DefaultKafkaProducerFactory factory = new DefaultKafkaProducerFactory<>(producerProperties());
        return new TracerKafkaProducerFactory<>(factory, kafkaSendTracer);
    }

    @Bean
    public KafkaTemplate<Integer, String> kafkaTemplate() {
        KafkaTemplate kafkaTemplate = new KafkaTemplate<Integer, String>(producerFactory(), true);
        kafkaTemplate.setDefaultTopic("default");
        return kafkaTemplate;
    }

    @Bean
    public KafkaListenerContainerFactory<ConcurrentMessageListenerContainer<Integer, String>> kafkaListenerContainerFactory() {
        ConcurrentKafkaListenerContainerFactory<Integer, String> factory = new ConcurrentKafkaListenerContainerFactory();
        factory.setConsumerFactory(consumerFactory());
        factory.setConcurrency(3);
        factory.getContainerProperties().setPollTimeout(3000);
        return factory;
    }


    @Bean
    public ConsumerFactory consumerFactory() {
        TraceSamplerConfig traceSamplerConfig = SpringContextUtils.getBean(TraceSamplerConfig.class);
        AbstractTracer kafkaConsumeTracer = KafkaConsumeTracer.getKafkaConsumeTracer(traceSamplerConfig);
        return new TracerKafkaConsumerFactory(new DefaultKafkaConsumerFactory<>(consumerProperties()), kafkaConsumeTracer);
    }

    @Bean
    @ConditionalOnMissingBean
    public KafkaListenerTracerAspect kafkaListenerTracerAspect() {
        TraceSamplerConfig traceSamplerConfig = SpringContextUtils.getBean(TraceSamplerConfig.class);
        AbstractTracer kafkaConsumeTracer = KafkaConsumeTracer.getKafkaConsumeTracer(traceSamplerConfig);
        return new KafkaListenerTracerAspect(kafkaConsumeTracer);
    }


    public Map<String, Object> producerProperties() {
        Map<String, Object> props = new HashMap<String, Object>();
        String retries = env.getProperty("kafka.producer.retries");
        String batchSize = env.getProperty("kafka.producer.batch.size");
        String lingerMs = env.getProperty("kafka.producer.linger.ms");
        String bufferMemory = env.getProperty("kafka.producer.buffer.memory");

        props.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, env.getProperty("kafka.producer.bootstrap.servers"));
        props.put(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, env.getProperty("kafka.producer.key.serializer"));
        props.put(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, env.getProperty("kafka.producer.value.serializer"));
        props.put(ProducerConfig.RETRIES_CONFIG, Integer.parseInt(retries));
        props.put(ProducerConfig.BATCH_SIZE_CONFIG, Integer.parseInt(batchSize));
        props.put(ProducerConfig.LINGER_MS_CONFIG, Integer.parseInt(lingerMs));
        props.put(ProducerConfig.BUFFER_MEMORY_CONFIG, Long.parseLong(bufferMemory));
        props.put(ProducerConfig.ACKS_CONFIG, env.getProperty("kafka.producer.acks", "all"));
        return props;
    }


    public Map<String, Object> consumerProperties() {
        Map<String, Object> propsMap = new HashMap<>();
        String enableAutoCommit = env.getProperty("kafka.consumer.enable.auto.commit");
        propsMap.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, env.getProperty("kafka.consumer.servers"));
        propsMap.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, Boolean.valueOf(enableAutoCommit));
        propsMap.put(ConsumerConfig.AUTO_COMMIT_INTERVAL_MS_CONFIG, env.getProperty("kafka.consumer.auto.commit.interval"));
        propsMap.put(ConsumerConfig.SESSION_TIMEOUT_MS_CONFIG, env.getProperty("kafka.consumer.session.timeout"));
        propsMap.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class);
        propsMap.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class);
        propsMap.put(ConsumerConfig.GROUP_ID_CONFIG, env.getProperty("kafka.consumer.group.id"));
        propsMap.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, env.getProperty("kafka.consumer.auto.offset.reset"));

        return propsMap;
    }

    @Override
    public void setEnvironment(Environment environment) {
        this.env = environment;
        SimpleTracerConfiguration.setApplicationName(getApplicationName());
    }

    private String getApplicationName() {
        String applicationName = this.env.getProperty("spring.application.name");
        return applicationName;
    }
}
